import onnx
import onnxruntime
import numpy as np
import torch
from torchvision.transforms import transforms
import h5py
from PIL import Image
import time

# model = onnx.load("BiFNet_32_44186.onnx")
# onnx.checker.check_model(model)
#
# session = onnxruntime.InferenceSession('BiFNet_32_44186.onnx')
# input_names = [inp.name for inp in session.get_inputs()]
# print(input_names)
# class_images = np.random.randn(2, 1, 224, 224).astype(np.float32)
# defect_images = np.random.randn(2, 1, 224, 224).astype(np.float32)
# inputs_dict = {
#     'class_images': class_images,
#     'defect_images': defect_images
# }
# outputs = session.run(None, inputs_dict)
# print(outputs[0].shape)



def predictImages(image, onnx_weight):

    DEVICE = torch.device("cpu")

    # preprocess
    transform_test = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            # transforms.Normalize(mean=[0.42368844, 0.42368844, 0.42368844], std=[0.14975642, 0.14975642, 0.14975642])
            transforms.Normalize(mean=[0.40272543], std=[0.13901867])
        ])


    image = Image.open(image).convert("RGB")
    image = transform_test(image)
    image = torch.unsqueeze(image, dim=0).to(DEVICE).numpy()

    ort_inputs = {"input": image}

    # Predict
    session = onnxruntime.InferenceSession(onnx_weight)
    start_time = time.time()
    for i in range(200):
        outputs = session.run(None, ort_inputs)
    end_time = time.time()
    print((end_time - start_time) / 200)
    return

if __name__ == '__main__':
    image = 'img.png'
    onnx_weight = 'convnext_s_224.onnx'
    predictImages(image, onnx_weight)

    '''
    onnx cpu
    convnext_s 224 0.053655390739440915
    swinV2_s 256 0.07726752638816833
    '''
